
import time
import numpy as np
from coppeliasim_zmqremoteapi_client import RemoteAPIClient
import os
import math
import basedriver
import roboticstoolbox as rtb
from motion_planning import Motion_Planning
from spatialmath import *
from armkinematics_advanced import Jaka,get_trajectory,get_object_pose

# 初始化手指关节角度（以度为单位）
left_wheel_speed, right_wheel_speed = basedriver.inverse(0, 0)

# 定义手指和机械臂关节名称列表
finger_joint_names = ['finger01_joint01', 'finger01_joint02', 'finger01_joint03', 'finger01_joint04']
arm_joint_names = ['Joint01', 'Joint02', 'Joint03', 'Joint04', 'Joint05', 'Joint06']

# 所有关节名称列表
all_joint_names = [
    *[f'finger0{i}_joint0{j}' for i in range(1, 5) for j in range(1, 5)],
    'LM_motor', 'RM_motor', 'Joint06', 'Joint05', 'Joint04', 'Joint03', 'Joint02', 'Joint01'
]

# 所有关节句柄字典
joint_handles = {}

def initialize_simulation():
    """初始化仿真环境"""
    client = RemoteAPIClient()
    sim = client.require('sim')
    sim.stopSimulation()
    print('成功与仿真软件建立通信！', flush=True)
    
    sim.loadScene(scene_path)
    print('成功加载场景！', flush=True)
    for joint_name in all_joint_names:
        handle = sim.getObject('/'+joint_name)
        joint_handles[joint_name] = handle
    sim.setStepping(True)
    print('开始仿真！', flush=True)
    sim.startSimulation()
    sim.setJointTargetVelocity(joint_handles['LM_motor'], 0)
    sim.setJointTargetVelocity(joint_handles['RM_motor'], 0)
    return sim

def get_initial_angles(sim, joint_names):
    """获取关节初始角度"""
    initial_angles_deg = []
    for joint_name in joint_names:
        current_angle = sim.getJointPosition(joint_handles[joint_name])
        current_angle_deg = math.degrees(current_angle)
        initial_angles_deg.append(current_angle_deg)

    return initial_angles_deg

def generate_sequence(start=10, end=-40):

    result = []
    for third in range(start, end-1, -1):
        fourth = -30 - third
        row = [0, 0, third, fourth]
        result.append(row)
    
    return result

def execute_arm_trajectory(sim, arm_trajectory):
    """执行机械臂轨迹"""
    for point_index, point in enumerate(arm_trajectory.q):
        for joint_index, joint_name in enumerate(arm_joint_names):
            sim.setJointTargetPosition(joint_handles[joint_name], point[joint_index])
        while True:
            if all(np.isclose(sim.getJointPosition(joint_handles[joint_name]), point[joint_index], atol=1e-2) 
                   for joint_index, joint_name in enumerate(arm_joint_names)):
                break
            sim.step()

def plan_and_execute_finger_trajectory(sim):
    finger_trajectory = np.radians(generate_sequence())
    finger_position = []
    for trajectory_point in range(50):
        for finger_index in range(1, 5):
            for joint_index in range(4):
                joint_name = f'finger0{finger_index}_joint0{joint_index+1}'
                sim.setJointTargetPosition(joint_handles[joint_name], finger_trajectory[trajectory_point][joint_index])
                
        while True:
            if all(np.isclose(sim.getJointPosition(joint_handles[joint_name]), finger_trajectory[trajectory_point][joint_index], atol=1e-3) 
                   for joint_index, joint_name in enumerate(finger_joint_names)):
                break
            sim.step()


        sim.setJointTargetVelocity(joint_handles['Joint06'], 0)
        sim.setJointTargetVelocity(joint_handles['LM_motor'], 0)
        sim.setJointTargetVelocity(joint_handles['RM_motor'], 0)

        forcevector1 ,__ = read_force_sensor_data(sim,"finger01_tip")
        forcevector2 ,__ = read_force_sensor_data(sim,"finger02_tip")
        forcevector3 ,__ = read_force_sensor_data(sim,"finger03_tip")
        forcevector4 ,__ = read_force_sensor_data(sim,"finger04_tip")

        if forcevector1[0] > 200 or forcevector2[0] > 200 or forcevector3[0] > 200 or forcevector4[0] > 200:
            print("力传感器检测到碰撞，停止手指运动！", flush=True)
            break

        sim.step()

    print("手指轨迹执行完毕！", flush=True)


def print_format_transformation_matrix(sim,object_name):

    object_handle = sim.getObject(object_name)
    matrix = sim.getObjectMatrix(object_handle)
    # 初始化一个4x4的单位矩阵
    transformation_matrix = [
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1]
    ]
    
    # 将输入的12元素矩阵填充到4x4矩阵中
    for i in range(3):
        for j in range(4):
            if j == 3:
                transformation_matrix[i][j] = matrix[i * 4 + j] * 1000
            else:
                transformation_matrix[i][j] = matrix[i * 4 + j]

    for row in transformation_matrix:
        print(' '.join(
            [f"{elem:.3e}" if len(str(elem).replace('.', '').replace('e', '').replace('+', '').replace('-', '')) > 4 else str(elem) for elem in row]
        ), flush=True)
    
    return transformation_matrix


def format_transformation_matrix(matrix):
    # 初始化一个4x4的单位矩阵
    transformation_matrix = [
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1]
    ]
    
    # 将输入的12元素矩阵填充到4x4矩阵中
    for i in range(3):
        for j in range(4):
            transformation_matrix[i][j] = matrix[i * 4 + j]
    
    return transformation_matrix

def read_force_sensor_data(sim,force_sensor_name):
    """读取力传感器数据"""
    force_sensor_handle = sim.getObject('/'+force_sensor_name)
    result, forceVector, torqueVector = sim.readForceSensor(force_sensor_handle)
    # if result:
    #     print("力传感器"+force_sensor_name+"数据：", flush=True)
    #     print("力向量：", forceVector, flush=True)
    #     print("力矩向量：", torqueVector, flush=True)
    # else:  
    #     print("读取力传感器数据失败！", flush=True)
    #     return None, None
    
    if  not result:
        print("读取力传感器数据失败！", flush=True)
        return None, None
    # 处理力传感器数据
    return forceVector, torqueVector




def main():
    sim = initialize_simulation()
    try:
        intial_angles = get_initial_angles(sim, arm_joint_names)
        print("机械臂初始角度:", intial_angles, flush=True)

        arm_trajectory1,arm_trajectory2 ,arm_trajectory3 ,arm_trajectory4= get_trajectory(get_object_pose(sim),intial_angles)
        execute_arm_trajectory(sim, arm_trajectory1)
        print("到达预抓取姿态")

        matrix2 =sim.getObjectMatrix(sim.getObject("/hand_respondable"))
        Hand_T= np.array([
            [matrix2[0], matrix2[1], matrix2[2], matrix2[3]*1000],
            [matrix2[4], matrix2[5], matrix2[6], matrix2[7]*1000],
            [matrix2[8], matrix2[9], matrix2[10], matrix2[11]*1000],
            [0, 0, 0, 1]
        ])

        print('Hand_T:', Hand_T,flush=True)
        
        execute_arm_trajectory(sim, arm_trajectory2)
        print("到达抓取姿态",flush=True)


        plan_and_execute_finger_trajectory(sim)
        print("抓取完成！", flush=True)

        read_force_sensor_data(sim,"finger01_tip")
        read_force_sensor_data(sim,"finger02_tip")
        read_force_sensor_data(sim,"finger03_tip")
        read_force_sensor_data(sim,"finger04_tip")

        execute_arm_trajectory(sim, arm_trajectory3)
        print("到达放置姿态")        

        execute_arm_trajectory(sim, arm_trajectory4)

        while True:
            sim.step()
    except KeyboardInterrupt:
        print('程序被用户终止！')
    finally:
        sim.stopSimulation()
        print('仿真已停止。')

if __name__ == "__main__":
    main()